package org.robotfish.learn.tests;

import static org.junit.Assert.*;

import java.util.ArrayList;
import java.util.Arrays;

import org.ejml.data.DenseMatrix64F;
import org.ejml.ops.CommonOps;
import org.ejml.ops.MatrixFeatures;
import org.ejml.simple.SimpleMatrix;
import org.junit.Test;
import org.robotfish.learn.Features;
import org.robotfish.learn.GaussianFeatures;
import org.robotfish.learn.GaussianProdFeatures;
import org.robotfish.learn.LinearModel;
import org.robotfish.learn.LinearRegression;
import org.robotfish.learn.PolyFeatures;

public class FeaturesTest {

	@Test
	public void testGaussian() {
		GaussianFeatures g1 = new GaussianFeatures(new double[]{0.1, 0.2, 0.5}, 0.5, 0);
		GaussianFeatures g2 = new GaussianFeatures(new double[]{0.1, 0.2, 0.5}, 2.5, 1);
		double[] data = {
			0.0, 0.1,
			0.2, 0.3,
			0.4, 0.7
		};
		SimpleMatrix m = new SimpleMatrix(3, 2, true, data);
		SimpleMatrix f1 = g1.featMatrix(m);
		assertEquals("g1(0)", 0.9607894, f1.get(0, 0), 1e-5);
		assertEquals("g1(1)", 0.9607894, f1.get(1, 0), 1e-5);
		assertEquals("g1(2)", 0.6976763, f1.get(2, 0), 1e-5);
		
		SimpleMatrix f2 = g2.featMatrix(m);
		assertEquals("g2(0)", 0.9747249, f2.get(0, 2), 1e-5);
		assertEquals("g2(1)", 0.9936204, f2.get(1, 2), 1e-5);
		assertEquals("g2(2)", 0.9936204, f2.get(2, 2), 1e-5);

	}
	
	@Test
	public void testCycleGaussian() {
		GaussianFeatures g1 = new GaussianFeatures(new double[]{0.1, 0.4}, 0.05, 1, 0.12);
		double[] data = {
			0, 0.00,
			0, 0.25,
			0, 0.50,
			0, 0.75,
			0, 1.00
		};
		SimpleMatrix m = new SimpleMatrix(5, 2, true, data);
		SimpleMatrix f1 = g1.featMatrix(m);
		
		assertEquals("g1(0, 0)", 0.8521438, f1.get(0, 0), 1e-5);
		assertEquals("g1(1, 0)", 0.6976763, f1.get(1, 0), 1e-5);
		assertEquals("g1(2, 0)", 0.5272924, f1.get(2, 0), 1e-5);
		assertEquals("g1(3, 0)", 0.3678794, f1.get(3, 0), 1e-5);
		assertEquals("g1(4, 0)", 0.2369278, f1.get(4, 0), 1e-5);
		
		assertEquals("g1(0, 1)", 0.5272924, f1.get(0, 1), 1e-5);
		assertEquals("g1(1, 1)", 0.6976763, f1.get(1, 1), 1e-5);
		assertEquals("g1(2, 1)", 0.8521438, f1.get(2, 1), 1e-5);
		assertEquals("g1(3, 1)", 0.9607894, f1.get(3, 1), 1e-5);
		assertEquals("g1(4, 1)", 1.0      , f1.get(4, 1), 1e-5);
		
		//String[] featNames = g1.featNames(new String[]{"x1", "x2"});
	}

	@Test
	public void testCycleGaussianProd() {
		GaussianProdFeatures g1 = new GaussianProdFeatures(new double[]{0.1, 0.4}, 0.05, 1, 0.12, 0, 2);
		double[] data = {
			 2, 0.00,
			 3, 0.25,
			 5, 0.50,
			 7, 0.75,
			-2, 1.00
		};
		SimpleMatrix m = new SimpleMatrix(5, 2, true, data);
		SimpleMatrix f1 = g1.featMatrix(m);
		
		assertEquals("g1(0, 0)", 0.8521438*4,  f1.get(0, 0), 1e-5);
		assertEquals("g1(1, 0)", 0.6976763*9,  f1.get(1, 0), 1e-5);
		assertEquals("g1(2, 0)", 0.5272924*25, f1.get(2, 0), 1e-5);
		assertEquals("g1(3, 0)", 0.3678794*49, f1.get(3, 0), 1e-5);
		assertEquals("g1(4, 0)", 0.2369278*4,  f1.get(4, 0), 1e-5);
		
		assertEquals("g1(0, 1)", 0.5272924*4,  f1.get(0, 1), 1e-5);
		assertEquals("g1(1, 1)", 0.6976763*9,  f1.get(1, 1), 1e-5);
		assertEquals("g1(2, 1)", 0.8521438*25, f1.get(2, 1), 1e-5);
		assertEquals("g1(3, 1)", 0.9607894*49, f1.get(3, 1), 1e-5);
		assertEquals("g1(4, 1)", 1.0*4      ,  f1.get(4, 1), 1e-5);
		
		//String[] featNames = g1.featNames(new String[]{"x1", "x2"});
	}

	@Test
	public void testPoly() {
		PolyFeatures f = new PolyFeatures(2, new int[]{1,3}, true, false);
		double[] data = {
			1, 0.00,
			2, 0.25,
			10, 0.50,
			14, 0.75,
			15, 1.00
		};
		
		SimpleMatrix m = new SimpleMatrix(5, 2, true, data);
		SimpleMatrix feat = f.featMatrix(m);
		
		assertEquals("rows", m.numRows(), feat.numRows());
		assertEquals("cols", 5, feat.numCols());
		double[] expected = {
			1.0, 1.0, 1.0, 0, 0,
			1.0, 2.0, 8.0, 0.25, Math.pow(0.25, 3),
			1.0, 10, 1000, 0.5, Math.pow(0.5, 3),
			1.0, 14, Math.pow(14, 3), 0.75, Math.pow(0.75, 3),
			1.0, 15, Math.pow(15, 3), 1.00, Math.pow(1.00, 3)
		};
		SimpleMatrix featExpected = new SimpleMatrix(feat.numRows(), feat.numCols(), true, expected);
		assertTrue( feat.isIdentical(featExpected, 1e-6) );
		
		// the same features without the constant term:
		f = new PolyFeatures(2, new int[]{1,3}, false, false);
		feat = f.featMatrix(m);
		// submatrix without the first column:
		SimpleMatrix featExpected2 = featExpected.extractMatrix(0, SimpleMatrix.END, 1, SimpleMatrix.END);
		assertTrue( feat.isIdentical(featExpected2, 1e-6) );
	}

	@Test
	public void testPolySign() {
		PolyFeatures f = new PolyFeatures(2, new int[]{1,2,3}, true, true);
		double[] data = {
			-1, 0.00,
			2, -0.25,
		};
		
		SimpleMatrix m = new SimpleMatrix(2, 2, true, data);
		SimpleMatrix feat = f.featMatrix(m);
		
		assertEquals("rows", m.numRows(), feat.numRows());
		assertEquals("cols", 7, feat.numCols());
		double[] expected = {
			1.0, -1.0, -1.0, -1.0, 0, 0, 0,
			1.0,  2.0,  4.0,  8.0, -0.25, -Math.pow(0.25, 2), Math.pow(-0.25, 3)
		};
		SimpleMatrix featExpected = new SimpleMatrix(feat.numRows(), feat.numCols(), true, expected);
		assertTrue( feat.isIdentical(featExpected, 1e-6) );
	}
	
	@Test
	public void binFunctions() {
		assertEquals( 0, LinearModel.binlog(0) );
		assertEquals( 0, LinearModel.binlog(1) );  // 2^0 = 1
		assertEquals( 1, LinearModel.binlog(2) );  // 2^1 = 2
		assertEquals( 2, LinearModel.binlog(4) );  // 2^2 = 4
		assertEquals( 3, LinearModel.binlog(15) );
		assertEquals( 4, LinearModel.binlog(16) ); 
		assertEquals( 4, LinearModel.binlog(31) ); 
		assertEquals( 5, LinearModel.binlog(32) ); // 2^5 = 32
		
		assertEquals( true, LinearModel.isPowerOf2(1) );
		assertEquals( true, LinearModel.isPowerOf2(2) );
		assertEquals( true, LinearModel.isPowerOf2(4) );
		assertEquals( true, LinearModel.isPowerOf2(8) );
		assertEquals( false, LinearModel.isPowerOf2(3) );
		assertEquals( false, LinearModel.isPowerOf2(7) );
		assertEquals( false, LinearModel.isPowerOf2(9) );
	}
	
	@Test
	public void matrixGrowing() {
		SimpleMatrix m = new SimpleMatrix(0, 3);
		for (int i=0; i<5; i++) {
			LinearModel.addRow(m, new SimpleMatrix( new double[][]{{0+i*3, 1+i*3, 2+i*3}} ));
		}
		assertEquals( 5, m.numRows() );
		assertEquals( 3, m.numCols() );
		SimpleMatrix mExpected = new SimpleMatrix( new double[][] {
			{0, 1, 2}, {3, 4, 5}, {6, 7, 8}, {9, 10, 11}, {12, 13, 14}
		});
		assertTrue( m.isIdentical(mExpected, 1e-7) );
		
		// speed test:
		m = new SimpleMatrix(0, 3);
		int n = 8000;
		for (int i=0; i<n; i++) {
			LinearModel.addRow( m, new double[]{0+i*3, 1+i*3, 2+i*3} );
		}
		assertEquals( n, m.numRows() );
		assertEquals( 3, m.numCols() );
	}
	
	@Test
	public void linearModelFeatures() {
		LinearModel lm = new LinearModel(3, 1.0);
		Features f1 = new PolyFeatures(new int[]{0, 1}, new int[]{1,3}, true, false);
		Features f2 = new GaussianFeatures(new double[]{0, 0.5}, 0.5, 1);
		lm.addFeature( f1 );
		lm.addFeature( f2 );
		
		SimpleMatrix fAll = lm.getFeatures( new SimpleMatrix(new double[][]{{10, 1.5, 20}}) );
		assertEquals( fAll.numCols(), f1.dimensions()+f2.dimensions() );
	}
	
	@Test
	public void variance() {
		double[] data = {
			-1, 0.00, 0, 1,
			2, 0.25, 0, 1,
			10, 0.50,0, 1,
			14, 0.75,0, 1,
			15, 1.00,0, 1
		};
			
		SimpleMatrix m = new SimpleMatrix(5, 4, true, data);
		double[] var = LinearModel.getColumnVariances( m );
		assertArrayEquals(new double[] {51.50000, 0.15625, 0.00000, 0.00000}, var, 1e-6);
	}
	
	@Test
	public void linearModelRegression() {
		// input features:
		double[] data = {
			-1, 0.00, 0.0, 1,
			2,  0.15, 0.1, 1,
			10, 0.50, 0.9, 1,
			14, 0.75, 0.2, 1,
			15, 1.00, 0.3, 1
		};
		SimpleMatrix X = new SimpleMatrix(5, 4, true, data);
		// output values:
		double[] dy = {
			15,
			30,
			2,
			3,
			-1
		};
		SimpleMatrix y = new SimpleMatrix(5, 1, true, dy);
		// linear regression parameters:
		SimpleMatrix alpha = LinearModel.linearRegression(X, y, new double[]{0.1, 0.2, 0.2, 0.0} ).alpha;
		SimpleMatrix alphaExpected = new SimpleMatrix(4, 1, true, new double[]{-1.2980263, -0.8034209, -5.1460207, 22.1136588});
		assertTrue( alpha.isIdentical(alphaExpected, 1e-7) );
	}
	
	@Test
	public void testUpdateReg() {
		// input features:
		double[] data = {
			-1, 0.00, 0.0, 1,
			2,  0.15, 0.1, 1,
			10, 0.50, 0.9, 1,
			14, 0.75, 0.2, 1,
			15, 1.00, 0.3, 1
		};
		SimpleMatrix X = new SimpleMatrix(5, 4, true, data);
		// output values:
		double[] dy = {
			15,	30,	2, 3, -1
		};
		SimpleMatrix y = new SimpleMatrix(5, 1, true, dy);
		// linear regression parameters:
		LinearModel lm = new LinearModel(4, 2.5);
		for (int i=0; i<dy.length; i++) {
			lm.addDataPoint( X.extractVector(true, i), y.get(i) );
		}
		lm.updateRegularization();
		assertEquals( 0, lm.getRegComponents().length );
		
		// linear regression parameters:
		lm = new LinearModel(4, 2.5);
		lm.addFeature( new PolyFeatures(4, new int[]{1}, false, false));
		for (int i=0; i<dy.length; i++) {
			lm.addDataPoint( X.extractVector(true, i), y.get(i) );
		}
		lm.updateRegularization();
		assertArrayEquals(new double[]{128.75, 0.426875, 0.3125, 0.0}, lm.getRegComponents(), 1e-7);
		
		lm.doFullRegression();
	}
	
	SimpleMatrix getData7X() {
		double[] data = {
				-1, 0.00, 0.0, 1,
				2,  0.15, 0.1, 1,
				10, 0.50, 0.9, 1,
				14, 0.75, 0.2, 1,
				15, 1.00, 0.3, 1,
				20, 1.75, 1.9, 1,
				 5, 2.00, 1.4, 1
			};
		SimpleMatrix X = new SimpleMatrix(7, 4, true, data);
		return X;
	}
	
	SimpleMatrix getData7y() {
		double[] dy = {
				15,	30,	2, 3, -1, 10, 5
			};
		SimpleMatrix y = new SimpleMatrix(dy.length, 1, true, dy);
		return y;
	}
	
	double[] getData7reg() {
		return new double[]{120.625, 0.28750000000000003, 0.4166666666666668, 0.0};
	}
	
	SimpleMatrix getData7alpha() {
		double[] alpha = {-0.474860, -5.906850,  3.271989, 16.498211};
		return new SimpleMatrix(alpha.length, 1, true, alpha);
	}
	
	@Test
	public void testIncrementalUpdate() {
		// data:
		SimpleMatrix X = getData7X();
		SimpleMatrix y = getData7y();
		double[] reg = getData7reg();
		
		// linear regression parameters:
		LinearModel lm = new LinearModel(4, 2.5);
		lm.addFeature( new PolyFeatures(4, new int[]{1}, false, false));
		
		// full regression, regularization is based on first 4 data points:
		lm.addDataPoints(X, y);
		lm.setRegComponents(reg);
		lm.doFullRegression();
		
		LinearRegression regression = lm.getRegression();
		// expected result:
//		regression.alpha.print();
		assertTrue( regression.alpha.isIdentical(getData7alpha(), 1e-5));
		
		// incremental regression:
		LinearModel lm2 = new LinearModel(4, 2.5);
		
		// full regression, regularization is based on first 4 data points:
		lm2.addFeature( new PolyFeatures(4, new int[]{1}, false, false) );
		int nfirst=4;
		lm2.addDataPoints(
			X.extractMatrix(0, nfirst, 0, SimpleMatrix.END), 
			y.extractMatrix(0, nfirst, 0, SimpleMatrix.END)
		);
		lm2.setRegComponents(reg);
		lm2.doFullRegression();
		SimpleMatrix alpha4 = new SimpleMatrix( new double[][]{ {-0.4372565}, {-8.5339734}, {-6.5107864}, {20.1729799} } );
		assertTrue(lm2.getRegression().alpha.isIdentical(alpha4, 1e-5));

		// checking if incremental updates give us the same answer: 
		for (int i=nfirst; i<X.numRows(); i++) {
			lm2.addPointAndUpdateRegression( X.extractVector(true, i), y.get(i) );
		}
		assertTrue( lm2.getRegression().alpha.isIdentical(getData7alpha(), 1e-5) );
		
		// checking also if predict works:
		assertEquals("prediction", 152.5004, lm2.getPrediction( new double[]{ 1.5, 2.8, 6.5, 9.0 }), 1e-3);
	}

}
